
#include "lcm_tiling.h"
#include "register/op_def_registry.h"
#include <algorithm>
using namespace std;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

	LcmTilingData tiling;
	
	// 获取输入tensor的形状
	auto shape_x1 = context->GetInputShape(0)->GetOriginShape();
	// 有多少个维度
	int dimNum1 = shape_x1.GetDimNum();
	int totalLength1 = 1;
	int dim1[3] = {0};
	// 计算每个维度的大小
	for (int i = 0; i < dimNum1; i++) {
		dim1[i] = shape_x1.GetDim(i);
		totalLength1 *= dim1[i];
	}
	tiling.set_totalLength1(totalLength1);

	// 获取输入tensor的形状
	auto shape_x2 = context->GetInputShape(1)->GetOriginShape();
	// 有多少个维度
	int dimNum2 = shape_x2.GetDimNum();
	int totalLength2 = 1;
	int dim2[3] = {0};
	// 计算每个维度的大小
	for (int i = 0; i < dimNum2; i++) {
		dim2[i] = shape_x2.GetDim(i);
		totalLength2 *= dim2[i];
	}
	tiling.set_totalLength2(totalLength2);

	tiling.set_dimNum1(dimNum1);
	tiling.set_dimNum2(dimNum2);
	tiling.set_dim1(dim1);
	tiling.set_dim2(dim2);

	int totalLength = 1;
	int dimNum = max(dimNum1, dimNum2);
	// printf("dimNum = %d\n", dimNum);
	int dim[3] = {0};
	for (int i = 0; i < dimNum; i++) {
		if (i < dimNum1) {
			dim[i] = max(dim[i], dim1[i]);
		}
		if (i < dimNum2) {
			dim[i] = max(dim[i], dim2[i]);
		}
		// printf("dim[%d] = %d\n", i, dim[i]);
		totalLength *= dim[i];
	}
	// printf("totalLength = %d\n", totalLength);
	tiling.set_totalLength(totalLength);
	tiling.set_dim(dim);
	tiling.set_dimNum(dimNum);
	
    // 获取输入tensor的数据类型
    auto tensor = context->GetInputTensor(0);
    auto type = tensor->GetDataType();
	if (totalLength == totalLength2) {
        context->SetTilingKey(1);
		// 40个核
		context->SetBlockDim(40);
	} else if (dimNum == 3 && totalLength == totalLength2 * dim[1] && type == ge::DT_INT64) {
        context->SetTilingKey(2);
		// 每个核处理最后一维
		context->SetBlockDim(totalLength / dim[dimNum - 1]);
	} else {
        context->SetTilingKey(3);
		// 每个核处理最后一维
		context->SetBlockDim(totalLength / dim[dimNum - 1]);
	}

	tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
	context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

	return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
	const gert::Shape* x1_shape = context->GetInputShape(0);
	gert::Shape* y_shape = context->GetOutputShape(0);
	*y_shape = *x1_shape;
	return GRAPH_SUCCESS;
}
static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
{
const auto inputDataType = context->GetInputDataType(0);
context->SetOutputDataType(0, inputDataType);
return ge::GRAPH_SUCCESS;
}
}


namespace ops {
class Lcm : public OpDef {
public:
	explicit Lcm(const char* name) : OpDef(name)
	{
		this->Input("input")
			.ParamType(REQUIRED)
			.DataType({ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
			.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
			.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
		this->Input("other")
			.ParamType(REQUIRED)
			.DataType({ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
			.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
			.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
		this->Output("out")
			.ParamType(REQUIRED)
			.DataType({ge::DT_INT8, ge::DT_INT16, ge::DT_INT32, ge::DT_INT64})
			.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
			.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

		this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

		this->AICore()
			.SetTiling(optiling::TilingFunc);
		this->AICore().AddConfig("ascend910b");

	}
};

OP_ADD(Lcm);
}
